package org.example;

import org.apache.commons.math3.distribution.PoissonDistribution;
import tech.tablesaw.api.DoubleColumn;
import tech.tablesaw.api.IntColumn;
import tech.tablesaw.api.Table;
import tech.tablesaw.plotly.Plot;
import tech.tablesaw.plotly.components.*;
import tech.tablesaw.plotly.traces.ScatterTrace;

import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Arrays;

public class UControlChart {

    public static void main(String[] args) throws Exception {
        int[] samples = Util.loadDataInt();
        Figure chart = getChart(samples);
        Plot.show(chart);
    }
    public static Figure getChart(int[] samples) throws Exception {
        int sampleCount = samples.length;          // 样本数量
        double defectRate = 0.1;       // 基准缺陷率

        // 生成模拟数据 - 缺陷数服从泊松分布
//        int[] defects = generateDefectData(sampleCount, samples, defectRate);
        double[] uValues = calculateUValues(samples, constantArray(10, samples.length));

        // 创建数据表
        Table data = Table.create("U Chart Data")
                .addColumns(
                        IntColumn.create("Sample", range(1, sampleCount + 1)),
                        IntColumn.create("Defects", samples),
                        IntColumn.create("SampleSize", constantArray(10, samples.length)),
                        DoubleColumn.create("U", uValues)
                );

        // 计算控制限
        double uBar = calculateUBar(samples, constantArray(10, samples.length));
        System.out.println("uBar:"+uBar);
        double[] ucl = new double[sampleCount];
        double[] lcl = new double[sampleCount];
        for (int i = 0; i < sampleCount; i++) {
            ucl[i] = uBar + 3 * Math.sqrt(uBar / 10);
            System.out.println("ucl:"+ucl[i]);

            lcl[i] = Math.max(0, uBar - 3 * Math.sqrt(uBar / 10)); // LCL不低于0
        }

        double uclBar = calculateBar(ucl);
        double lclBar = calculateBar(lcl);

        // 添加控制限列
        data.addColumns(
                DoubleColumn.create("UCL", constantArray(uclBar, sampleCount)),
                DoubleColumn.create("CL", constantArray(uBar, sampleCount)),
                DoubleColumn.create("LCL", constantArray(lclBar, sampleCount))
        );

        // 构建图表
        Figure uChart = createUChart(data);
        return uChart;
    }

    private static double calculateBar(double[] cls) {
        int length = cls.length;
        double t = 0.0;
        for (double cl:cls){
            t+=cl;
        }
        return t/length;
    }

    // 生成样本容量数组（可变样本大小）
    private static int[] generateSampleSizes(int count, int min, int max) {
        int[] sizes = new int[count];
        for (int i = 0; i < count; i++) {
            sizes[i] = min + (int)(Math.random() * (max - min + 1));
        }
        return sizes;
    }

    // 生成缺陷数据（泊松分布）
    private static int[] generateDefectData(int count, int[] sampleSizes, double rate) {
        int[] defects = new int[count];
        for (int i = 0; i < count; i++) {
            defects[i] = new PoissonDistribution(sampleSizes[i] * rate).sample();
        }
        return defects;
    }

    // 计算U值（单位缺陷数）
    private static double[] calculateUValues(int[] defects, int[] sampleSizes) {
        double[] u = new double[defects.length];
        for (int i = 0; i < defects.length; i++) {
            u[i] = (double) defects[i] / sampleSizes[i];
        }
        return u;
    }

    // 计算平均U值
    private static double calculateUBar(int[] defects, int[] sampleSizes) {
        int totalDefects = Arrays.stream(defects).sum();
        int totalSize = Arrays.stream(sampleSizes).sum();
        return (double) totalDefects / totalSize;
    }

    // 创建U控制图
    private static Figure createUChart(Table data) {
        Layout layout = Layout.builder()
                .title("U Control Chart (Defects per Unit)")
                .xAxis(Axis.builder().title("Sample Number").build())
                .yAxis(Axis.builder().title("Defects per Unit (U)").build())
                .width(1000)
                .height(600)
                .build();

        ScatterTrace dataTrace = ScatterTrace.builder(
                        data.intColumn("Sample"),
                        data.doubleColumn("U"))
                .name("U Value")
                .mode(ScatterTrace.Mode.LINE_AND_MARKERS)
                .build();

        ScatterTrace uclTrace = ScatterTrace.builder(
                        data.intColumn("Sample"),
                        data.doubleColumn("UCL"))
                .name("UCL")
                .mode(ScatterTrace.Mode.LINE)
                .line(Line.builder().color("red").dash(Line.Dash.DASH).width(1.5).build())
                .build();

        ScatterTrace clTrace = ScatterTrace.builder(
                        data.intColumn("Sample"),
                        data.doubleColumn("CL"))
                .name("CL")
                .mode(ScatterTrace.Mode.LINE)
                .line(Line.builder().color("green").dash(Line.Dash.DOT).width(1.5).build())
                .build();

        ScatterTrace lclTrace = ScatterTrace.builder(
                        data.intColumn("Sample"),
                        data.doubleColumn("LCL"))
                .name("LCL")
                .mode(ScatterTrace.Mode.LINE)
                .line(Line.builder().color("red").dash(Line.Dash.DASH).width(1.5).build())
                .build();

        return new Figure(layout, dataTrace, uclTrace, clTrace, lclTrace);
    }

    // 构建完整HTML
    private static String buildHtml(Figure figure, String title) {
        return "<!DOCTYPE html>\n" +
                "<html>\n" +
                "<head>\n" +
                "    <title>" + title + "</title>\n" +
                "    <script src='https://cdn.plot.ly/plotly-latest.min.js'></script>\n" +
                "</head>\n" +
                "<body>\n" +
                "    <div id='chart'></div>\n" +
                figure.asJavascript("chart") + "\n" +
                "</body>\n" +
                "</html>";
    }

    // 辅助方法
    private static int[] range(int start, int end) {
        return java.util.stream.IntStream.range(start, end).toArray();
    }
    private static double[] constantArray(double value, int length) {
        double[] arr = new double[length];
        Arrays.fill(arr, value);
        return arr;
    }
    private static int[] constantArray(int value, int length) {
        int[] arr = new int[length];
        Arrays.fill(arr, value);
        return arr;
    }
}